//===-     -----===//
//
//                     Damsl Lab LLVM-MPI Project
//
//
//===----------------------------------------------------------------------===//
//
//
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "mpicommpattern"
#include "llvm/Analysis/DebugInfo.h"
#include "llvm/Analysis/Passes.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/LoopPass.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CFG.h"
#include "llvm/Support/Debug.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/MPI/MPI.h"
#include "llvm/Pass.h"
#include "llvm/Constant.h"
#include "llvm/Module.h"
#include "llvm/Instructions.h"

using namespace llvm;

namespace {

class MPICommPattern : public FunctionPass {
public:
	enum CtrlStructType {
		IF = 0, // all branch that have > 1 successors
		FOR
	};
	struct CtrlStruct {
		CtrlStructType type;
		union {
			Loop* loop;
			BasicBlock* BB;
		} pointer;
		int pos;
	};

	static char ID; // Class identification, replacement for typeinfo
	MPICommPattern() : FunctionPass(ID) {}

private:
	LoopInfo *LI;
	typedef std::map<const Instruction *, std::vector<CtrlStruct>*>
	InstCtrlStuctSetMapTy;
	InstCtrlStuctSetMapTy result;
	//Value* traceup(Value* v);
	//Value* traceup_out_of_procudure(Argument* arg);
	//void tracedown(Value* v, SmallSet<Value*, 6> &result);
	bool isMpiCommCall(const Instruction &II);
	bool isMpiSendCall(const Instruction &II);
	inline bool funcname(const CallInst* II, const char* name) {
		Function* func = II->getCalledFunction();
		StringRef fname = func->getName();
		if (fname.startswith(StringRef(name))) {
			return true;
		} else {return false;}
	}
	inline bool funcname(const CallInst* II, StringRef name) {
		Function* func = II->getCalledFunction();
		StringRef fname = func->getName();
		if (fname.startswith(name)) {
			return true;
		} else {return false;}
	}
	void getAnalysisUsage(AnalysisUsage &AU) const;
	bool runOnFunction(Function &F);
	void IdentifyIF(std::vector<BasicBlock *> *DomBBs, Instruction *CurInst);
	~MPICommPattern();
};

// Register this pass...
char MPICommPattern::ID = 0;
static RegisterPass<MPICommPattern> X("mpicommpattern",
		"cache MPI communication pattern", false, true);

bool BBCmp(BasicBlock* BB1, BasicBlock* BB2) {
	return BB1->size() > BB2->size();
}


class BBDomCompare {
	const DominatorTree *const DT;
public:
	explicit BBDomCompare(const DominatorTree *dt) : DT(dt) {}
	bool operator()(const MPICommPattern::CtrlStruct LHS,
			const MPICommPattern::CtrlStruct RHS) const {
		return DT->dominates(
				LHS.pointer.BB->getTerminator()->getSuccessor(LHS.pos),
				RHS.pointer.BB->getTerminator()->getSuccessor(RHS.pos));
	}
};

}  // End of anonymous namespace


void MPICommPattern::getAnalysisUsage(AnalysisUsage &AU) const {
	AU.addRequired<LoopInfo>();
	AU.addRequired<DominatorTree>();
	AU.setPreservesAll();
}

void MPICommPattern::IdentifyIF(std::vector<BasicBlock *> *DomBBs,
		Instruction *CurInst) {
	DominatorTree *DT = &getAnalysis<DominatorTree>();
	std::vector<CtrlStruct> TempList;
	BasicBlock *ParentBB = CurInst->getParent();
	int position = 0;
	for(unsigned i = 0; i < DomBBs->size(); i++) {
		position = 0;
		for (succ_iterator SI = succ_begin((*DomBBs)[i]),
				E = succ_end((*DomBBs)[i]); SI != E; ++SI, ++position) {
			BasicBlock *Succ = *SI;
			if(DT->dominates(Succ, ParentBB)) {
				DEBUG(errs() << Succ->getName() << "\n");
				CtrlStruct CS;
				DEBUG(errs() << "if\n");
				CS.type = IF;
				CS.pos = position;
				CS.pointer.BB = (*DomBBs)[i];
				// sequence in this vector is significant.
				// will sort DominatorBBs later.
				TempList.push_back(CS);

				// only one of DomBBs[x]'s successor dominate ParentBB.
				break;
			}
		}
	}
	if(TempList.size() == 0) return;

	// sort TempList
	// 0 means outmost control structure.

	std::sort(TempList.begin(), TempList.end(), BBDomCompare(DT));
	if(result[CurInst] == 0) result[CurInst] = new std::vector<CtrlStruct>();
	result[CurInst]->insert(result[CurInst]->begin(),
			TempList.begin(), TempList.end());
}

// work for MPI_Send, MPI_Isend, MPI_Recv, MPI_Irecv
// consider if, switch, for loop, while loop;
// do while loop is not supported, continue, break is not supported.
bool MPICommPattern::runOnFunction(Function &F) {
	LI = &getAnalysis<LoopInfo>();

	std::vector<Instruction *> MpiCallList;

	// populate list with all mpi calls need to be processed.
	for (Function::iterator BI = F.begin(), BE = F.end(); BI
	!= BE; ++BI) {
		for (BasicBlock::iterator II = (*BI).begin() , E = (*BI).end(); II
		!= E; ++II) {
			//DEBUG(errs() << "iteration\n");
			if (isMpiCommCall(*II)) {
				MpiCallList.push_back(&*II);
			}
		}
	}

	// All basic blocks in this function
	std::vector<BasicBlock *> AllBBs;
	for (Function::iterator BI = F.begin(), BE = F.end(); BI
	!= BE; ++BI) {
		AllBBs.push_back(&*BI);
	}

	// All basic blocks in all loops
	std::vector<BasicBlock *> AllLoopBBs;
	for (LoopInfo::iterator I = LI->begin(), E = LI->end(); I != E; I++) {
		AllLoopBBs.insert(AllLoopBBs.begin(), (*I)->getBlocks().begin(),
				(*I)->getBlocks().end());
	}

	// OutmostBBs include basic blocks that does not belong to any loop
	std::vector<BasicBlock *> OutmostBBs;

//	std::sort(AllBBs.begin(), AllBBs.end(), BBCmp);
//	std::sort(AllLoopBBs.begin(), AllLoopBBs.end(), BBCmp);
//	std::set_difference(AllBBs.begin(), AllBBs.end(),
//			AllLoopBBs.begin(), AllLoopBBs.end(), back_inserter(OutmostBBs), BBCmp);

	std::sort(AllBBs.begin(), AllBBs.end());
	std::sort(AllLoopBBs.begin(), AllLoopBBs.end());
	std::set_difference(AllBBs.begin(), AllBBs.end(),
			AllLoopBBs.begin(), AllLoopBBs.end(), back_inserter(OutmostBBs));


	for(unsigned i = 0; i < AllBBs.size(); ++i) {
		DEBUG(errs() << AllBBs[i]->getName() << "  ");
	}
	DEBUG(errs() << "\n");
	for(unsigned i = 0; i < AllLoopBBs.size(); ++i) {
		DEBUG(errs() << AllLoopBBs[i]->getName() << "  ");
	}
	DEBUG(errs() << "\n");
	for(unsigned i = 0; i < OutmostBBs.size(); ++i) {
		DEBUG(errs() << OutmostBBs[i]->getName() << "  ");
	}
	DEBUG(errs() << "\n");

	// OutmostIfBBs include basic blocks that does not belong to any loop
	// and whose last instruction is a branch instruction having more
	// than one branching target or a switch instruction.
	std::vector<BasicBlock *> OutmostIfBBs;
	for(std::vector<BasicBlock *>::const_iterator VI = OutmostBBs.begin(),
			VE = OutmostBBs.end(); VI != VE; VI++) {
		if((*VI)->getTerminator()->getNumSuccessors() > 1)
			OutmostIfBBs.push_back(*VI);
	}

	// current mpi call instruction that is processed
	Instruction *CurInst;
	std::vector<BasicBlock *> LoopBBs;
	std::vector<BasicBlock *> SubLoopBBs;
	std::vector<BasicBlock *> OutBBs;
	std::vector<BasicBlock *> OutIfBBs;
	while(MpiCallList.empty() == false) {
		DEBUG(errs() << "process Instruction\n");
		CurInst = MpiCallList.front();
		MpiCallList.erase(MpiCallList.begin());
		result[CurInst] = 0;
		//std::vector<CtrlStruct>*  CtrlStructNest;
		// see if this instruction in the same basic block as some mpi call
		// that has been processed. If so, avoid recomputation.
		// use a bool to avoid using goto to jump over two nesting loop structure.
		//DEBUG(errs() << "  check skipcomputation\n");
		bool skipcomputation = false;
		for(InstCtrlStuctSetMapTy::const_iterator RI = result.begin(),
				RE = result.end(); RI != RE; RI++) {
			if(RI->first != CurInst &&
					RI->first->getParent() == CurInst->getParent()) {
				//DEBUG(errs() << "    found duplicate\n");
				if(RI->second != 0) {
					std::vector<CtrlStruct>* t = new std::vector<CtrlStruct>();
					//t->assign(RI->second->begin(), RI->second->end());
					*t = *(RI->second);
					result[CurInst] = t;
				}
				skipcomputation = true;
				break;
			}
		}
		if(skipcomputation) {
			DEBUG(errs() << "  skipcomputation\n");
			continue;
		}

		// process one mpi call
		BasicBlock *ParentBB = CurInst->getParent();
		Loop *enclosingloop = LI->getLoopFor(ParentBB);
		// FIXME: worthwhile to consider this uncommon case?
		// mpi call does not live in any loop

		// mpi call live in some loop
		// iterate up through nested loops

		for (Loop *CurLoop = enclosingloop; CurLoop;
										CurLoop = CurLoop->getParentLoop()) {
			LoopBBs.clear();
			SubLoopBBs.clear();
			OutBBs.clear();
			OutIfBBs.clear();
			if(result[CurInst] == 0) result[CurInst] = new std::vector<CtrlStruct>();
			// see if this loop contain if statement
			for(Loop::block_iterator BI = CurLoop->block_begin(),
					E = CurLoop->block_end(); BI != E; BI++) {
				if((*BI)->getTerminator()->getNumSuccessors() > 1 &&
						*BI != CurLoop->getHeader() &&
						*BI != CurLoop->getLoopLatch()) {
					LoopBBs.push_back(*BI);
				}
			}
			for(Loop::iterator LI = CurLoop->begin(),
											LE = CurLoop->end(); LI != LE; LI++) {
				for(Loop::block_iterator BI = (*LI)->block_begin(),
						BE = (*LI)->block_end(); BI != BE; BI++) {
					SubLoopBBs.push_back(*BI);
				}
			}

			std::sort(LoopBBs.begin(), LoopBBs.end());
			std::sort(SubLoopBBs.begin(), SubLoopBBs.end());
			std::set_difference(LoopBBs.begin(), LoopBBs.end(), SubLoopBBs.begin(),
														SubLoopBBs.end(), back_inserter(OutBBs));
			for(std::vector<BasicBlock *>::const_iterator VI = OutBBs.begin(),
					VE = OutBBs.end(); VI != VE; VI++) {
				if((*VI)->getTerminator()->getNumSuccessors() > 1)
					OutIfBBs.push_back(*VI);
			}
			IdentifyIF(&OutIfBBs, CurInst);

			CtrlStruct CS;
			CS.type = FOR;
			CS.pos = -1;
			CS.pointer.loop = CurLoop;
			result[CurInst]->insert(result[CurInst]->begin(), CS);
		}
		// add OutmostBBs that dominate mpi call(and its parent basic block)
		if(result[CurInst] != 0)
			DEBUG(errs() << result[CurInst]->size() <<"\n");
		//DEBUG(errs() << OutmostIfBBs.size() <<"\n");
		IdentifyIF(&OutmostIfBBs, CurInst);
		if(result[CurInst] != 0)
			DEBUG(errs() << result[CurInst]->size() <<"\n");
	}

	//debug printout
	DEBUG(errs() << "size of result is " << result.size() << "\n");
	for(InstCtrlStuctSetMapTy::const_iterator i = result.begin(),
				e = result.end(); i != e; i++) {
		std::vector<CtrlStruct>* m = std::pair<const Instruction *,
																		std::vector<CtrlStruct>*>(*i).second;
		if (m == 0) {DEBUG(errs() << 0 << "\n");}
		else {DEBUG(errs() << m->size() << "\n");}
		for(unsigned j = 0; j < m->size(); ++j) {
			DEBUG(errs() << "  " << ((*m)[j]).type << "\n");
		}
	}
	return false;
}

MPICommPattern::~MPICommPattern() {
	for(InstCtrlStuctSetMapTy::const_iterator i = result.begin(),
			e = result.end(); i != e; i++) {
		delete std::pair<const Instruction *, std::vector<CtrlStruct>*>(*i).second;
	}
}


// Return true if Instruction II is a MPI Communication call.
bool MPICommPattern::isMpiCommCall(const Instruction &II) {
	if (const CallInst* callInst = dyn_cast<CallInst>(&II)) {
		Function* func = callInst->getCalledFunction();
		StringRef fname = func->getName();
		if (fname.startswith(StringRef("MPI_Send"))  ||
				fname.startswith(StringRef("MPI_Isend")) ||
				fname.startswith(StringRef("MPI_Recv"))  ||
				fname.startswith(StringRef("MPI_Irecv")))
			return true;
	}
	return false;
}

// Return true if Instruction II is a MPI send call.
bool MPICommPattern::isMpiSendCall(const Instruction &II) {
	if (const CallInst* callInst = dyn_cast<CallInst>(&II)) {
		Function* func = callInst->getCalledFunction();
		StringRef fname = func->getName();
		if (fname.startswith(StringRef("MPI_Send"))  ||
				fname.startswith(StringRef("MPI_Isend")))
			return true;
	}
	return false;
}


